from ..config import hyper
from typing import List
from ..preprocess.seq2umt_loader import Seq2UmtDataset, Seq2umt_loader


class DataPreProcessor:
    def __init__(self, raw_data_list: List = None, use_dataset_file: bool = False, dataset_file=None):
        self.raw_data_list = raw_data_list
        self.data_loader = None
        self.use_dataset_file = use_dataset_file
        self.dataset_file = dataset_file

    def process(self):
        prepared_data = []
        if self.use_dataset_file:
            data_set = Seq2UmtDataset(hyper=hyper, dataset=self.dataset_file)
            batch_size = hyper.batch_size_eval
        else:
            for line in self.raw_data_list:
                if len(line) < hyper.max_text_len:
                    prepared_data.append(line)
            batch_size = 1 if len(prepared_data) <= 10 else min(200, len(prepared_data) // 10)
            data_set = Seq2UmtDataset(hyper=hyper, sentence_list=prepared_data)
        self.data_loader = Seq2umt_loader(
            data_set,
            batch_size=batch_size,
            pin_memory=True,
            num_workers=0,
        )
        return self

    def get_data_loader(self):
        if self.data_loader is None:
            self.process()
        return self.data_loader
